# coding: utf-8
"""
pip install sympy
"""
import numpy as np
from sympy import symbols, diff, solve, Eq

# 将函数f中的参数替换成实际数值,返回一个新的函数表达式
# 参数:
#   s -- 参数词典(s[a],s[b],s[c])
def __sub_expr__(expr, s):
    return expr.subs({"a": s[0], "b": s[1], "c": s[2]})


# 给定xy坐标,构造方程表达式(f(x)=y)
# 参数:
#   xy -- 坐标列表, 示例: [(10.3, 20), (11.5, 40.5), ... ]
def __build_exp__(xy, expr):
    exprs = []
    for x, y in xy:
        exprs.append(Eq(expr.subs({"x": x}), y))
    return exprs


# 已知(x1,y1) (x2,y2) (x3,y3)，求抛物线的参数a,b,c
# 方法:
#   解方程组：y1=f(x1);y2=f(x2);y3=f(x3), 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 3个点坐标, 示例: [(x1,y1), (x2,y2), (x3,y3)]
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_three_points(xy):
    # 声明符号变量
    a, b, c, x = symbols("a b c x")
    # 定义抛物线表达式
    expr = a * x ** 2 + b * x + c
    # 替换表达式中的x和y为具体数值
    exprs = __build_exp__(xy, expr)
    # 求解表达式
    s = solve(exprs, [a, b, c])
    return s[a], s[b], s[c]


# 已知2个点以及第1个点的导数求抛物线方程
# 方法:
#   解方程组：y1=f(x1);y2=f(x2);f'(x1)=k, 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 2个点坐标, 示例: [(x1,y1), (x2,y2)]
#   k -- 第1个点的导数
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_two_points_and_diff(xy, k, i=0):
    # 声明符号变量
    a, b, c, x = symbols("a b c x")
    # 定义抛物线表达式
    expr = a * x ** 2 + b * x + c
    # 定义抛物线一阶导表达式
    diff_expr = diff(expr, x)
    # 替换表达式中的x和y为具体数值
    exprs = __build_exp__(xy, expr)
    exprs.append(Eq(diff_expr.subs({"x": xy[i][0]}), k))
    # 求解表达式
    s = solve(exprs, [a, b, c])
    return s[a], s[b], s[c]


# 已知2个点以及第1点的导数求抛物线方程
# 方法:
#   解方程组：y1=f(x1);y2=f(x2);f'(x1)=k, 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 2个点坐标, 示例: [(x1,y1), (x2,y2)]
#   k -- 第1个点的导数
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_two_points_and_diff_1(xy, k):
    return solve_parabolic_two_points_and_diff(xy, k, i=0)


# 已知2个点以及第2点的导数求抛物线方程
# 方法:
#   解方程组：y1=f(x1);y2=f(x2);f'(x2)=k, 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 2个点坐标, 示例: [(x1,y1), (x2,y2)]
#   k -- 第2个点的导数
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_two_points_and_diff_2(xy, k):
    # 链表反转后再调用
    return solve_parabolic_two_points_and_diff(xy, k, i=1)


# 已知2个点以及2个点的导数求抛物线方程
# 方法:
#   解方程组：y1=f(x1);f'(x1)=k1;f'(x2)=k2, 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   xy -- 2个点坐标, 示例: [(x1,y1), (x2,y2)]
#   k -- 第2个点的导数
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
# 注1: 第2点的y值其实并没有用到, 任意设置即可
# 注2: 如果坐标点是极值点, 则导数为0!
def solve_parabolic_two_points_and_diff_3(xy, k1, k2):
    # 声明符号变量
    a, b, c, x = symbols("a b c x")
    # 定义抛物线表达式
    expr = a * x ** 2 + b * x + c
    # 定义抛物线一阶导表达式
    diff_expr = diff(expr, x)
    # 替换表达式中的x和y为具体数值
    exprs = __build_exp__(xy[:1], expr)
    exprs.append(diff_expr.evalf(subs={"x": xy[0][0]}) == k1)
    exprs.append(diff_expr.evalf(subs={"x": xy[1][0]}) == k2)
    # 求解表达式
    s = solve(exprs, [a, b, c])
    return s[a], s[b], s[c]


# 已知2条相切的抛物线, 求相切的交点
# 方法:
#   解方程组：f1(x)=f2(x); f1'(x)=f2'(x), 其中y=f(x)=a*x^2+b*x+c
# 参数:
#   s1  --  第1条抛物线的参数, 示例: -0.01, 16.5, 1805
#   s2  --  第2条抛物线的参数, 示例: -0.01, 16.5, 1805
# 返回值:
#   交点, 示例: { "x": 50, "y": 420.56 }
def solve_two_parabolics_jiaodian(s1, s2):
    # 声明符号变量
    a, b, c, x = symbols("a b c x")
    # 定义抛物线表达式
    expr = a * x ** 2 + b * x + c
    # 定义抛物线一阶导表达式
    diff_expr = diff(expr, x)
    # 构建方程表达式
    exprs = [
        # Eq(__sub_expr__(expr, s1), __sub_expr__(expr, s2)),
        Eq(__sub_expr__(diff_expr, s1), __sub_expr__(diff_expr, s2)),
    ]
    # 求解表达式
    s = solve(exprs, [x])
    return s[x], expr.evalf(subs={"x": s[x], "a": s1[0], "b": s1[1], "c": s1[2]})


# 已知第1条抛物线 s 和 第2条抛物线的极值点, 求第2条抛物线的参数
# 注: 两条抛物线相切
# 方法:
#   解方程组：f1'(x)=f2'(x); y2=f2(x); f2'(x_min)=0
# 参数:
#   s      -- 第一条抛物线参数, 示例: s = -0.0446064374845907, 9.81341624660994, -139.357893563547
#   xy_min -- 第2条抛物线的极值点坐标
# 返回值:
#   元组或列表, 示例: -0.0446064374845907, 9.81341624660994, -139.357893563547
def solve_parabolic_by_min_point(s, xy_min):
    # 声明符号变量
    a, b, c, x = symbols("a b c x")
    # 定义抛物线表达式
    expr = a * x ** 2 + b * x + c
    # 定义抛物线一阶导表达式
    diff_expr = diff(expr, x)
    # 替换表达式中的x和y为具体数值
    exprs = [
        Eq((s[1] - b) ** 2 - 4 * (s[0] - a) * (s[2] - c), 0),
        Eq(expr.subs({"x": xy_min[0]}), xy_min[1]),
        Eq(diff_expr.subs({"x": xy_min[0]}), 0),
    ]
    s2 = solve(exprs, [a, b, c], solution_dict=True)[0]
    return s2[0], s2[1], s2[2]


# 给定一组数据xy,共3个点坐标, 计算2条抛物线的参数及相切点
# 右抛物线(前2个): 工况点、极大工况点
# 左抛物线(第3个): 极小工况点
def solve_two_parabolics_1(xy):
    # 根据工况点和最大工况点求右抛物线
    s1 = solve_parabolic_two_points_and_diff_2(xy[:2], 0)
    # 根据右抛物线+左抛物线最小工况点 求 左抛物线
    s2 = solve_parabolic_by_min_point(s1, xy[2])
    # 两抛物线相切点
    p0 = solve_two_parabolics_jiaodian(s1, s2)
    return s1[0], s1[1], s1[2], s2[0], s2[1], s2[2], p0[0], p0[1]


# 给定一组X和Y数据, 进行多项式拟合
# 参数:
#   xy   -- 坐标点数据, 示例: [(x1, y1), (x2, y2), ... ]
#   n    -- 拟合次数
# 返回值:
#   列表, 示例: [-0.2506347664975213, 64.65774115325407, -1326.4090233350287]
# 注: 从高阶 -> 低阶, 例如: [a, b, c]
def polyfit(xy, n=2):
    return list(np.polyfit([x for x, _ in xy], [y for _, y in xy], n))


# 给定一组数据xy,共3个点坐标, 计算2条抛物线的参数及相切点
# 右抛物线(前2个): 工况点、极大工况点
# 左抛物线(第3个): 极小工况点
def solve_two_parabolics_2(xy):
    xk, yk = xy[0]
    xmax, ymax = xy[1]
    xmin, ymin = xy[2]
    # 求右抛物线
    a1 = (yk - ymax) / ((xk - xmax) ** 2 + 1.0e-12)
    b1 = -2 * xmax * (yk - ymax) / ((xk - xmax) ** 2 + 1.0e-12)
    c1 = xmax ** 2 * (yk - ymax) / ((xk - xmax) ** 2 + 1.0e-12) + ymax

    # 求左抛物线
    L1 = b1 ** 2 - 4 * a1 * c1 + 4 * a1 * ymin
    L2 = a1 * xmin ** 2 + b1 * xmin + c1 - ymin
    a2 = -1 * L1 / (4 * L2 + 1.0e-12)
    b2 = L1 * xmin / (2 * L2 + 1.0e-12)
    c2 = ymin - L1 * xmin ** 2 / (4 * L2 + 1.0e-12)

    # 求相切点
    x0 = -0.5 * (b1 - b2) / (a1 - a2 + 1.0e-12)
    y0 = a1 * x0 ** 2 + b1 * x0 + c1
    return a1, b1, c1, a2, b2, c2, x0, y0


def main():
    xy = [(110, 400.38), (151.19, 324.7)]
    # # 已知2点坐标, 且第1点为极值点(导数为0), 计算抛物线
    s1 = solve_parabolic_two_points_and_diff_1(xy, 0)
    print(s1)

    # 已知工况点(右抛物线)、极大工况点(右抛物线)、极小工况点(左抛物线), 求2条相切抛物线
    xy = [(151.19, 324.70), (110.00, 400.38), (45.00, 228.217)]
    a1, b1, c1, a2, b2, c2, x0, y0 = solve_two_parabolics_1(xy)
    print("左抛物线:", a2, b2, c2)
    print("右抛物线:", a1, b1, c1)
    print("左右抛物线相切点:", x0, y0)


if __name__ == "__main__":
    main()
